#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import argparse
import subprocess
import pandas as pd
import json
from datetime import datetime
import numpy as np
import paddle
import yaml
import sys

# 确保能导入a1模块
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from a1 import Classifier, TIMITDataset, load_config, get_experiment_config, same_seeds

# 定义要运行的实验
EXPERIMENTS = [
    "task1_todo1",    # 四层网络
    "task1_todo2",    # 使用ReLU激活函数
    "task1_todo3",    # 添加批归一化
    "task1_todo4",    # 增加训练轮数
    "task2_dropout",  # 增加Dropout
    "task2_learning_rate", # 调整学习率
    "task2_wider_net",  # 调整网络宽度
    "task2_l2_regularization",  # 添加L2正则化
    "task2_last_layer_softmax",  # 最后一层使用Softmax激活函数
    "task2_larger_batch_size"  # 增加batch size
]

def run_experiment(experiment=None, config_file="config.yaml"):
    """运行单个实验"""
    cmd = ["python", "a1.py", "--config", config_file]
    
    if experiment:
        cmd.extend(["--experiment", experiment])
        print(f"Running experiment: {experiment}")
    else:
        print("Running baseline experiment")
    
    # 执行命令
    process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    stdout, stderr = process.communicate()
    
    if process.returncode != 0:
        print(f"Error running experiment: {stderr}")
        return False
    
    print(f"Experiment completed successfully")
    return True

def inference_with_model(model_name, data_path=None, config_file="config.yaml", batch_size=64, output_file=None):
    """
    使用训练好的模型进行推理，并将结果保存到文件
    
    参数:
        model_name: 要使用的模型名称，对应config中的model_name
        data_path: 要推理的数据集路径，如果为None则使用config中的private_test_path
        config_file: 配置文件路径
        batch_size: 批处理大小
        output_file: 输出文件路径，如果为None则使用config中的prediction_file
    
    返回:
        predictions: 预测结果列表
    """
    print(f"正在使用模型 {model_name} 进行推理...")
    
    # 加载配置
    config = load_config(config_file)
    
    # 查找对应的实验配置
    model_config = None
    model_experiment = None
    
    # 首先检查基础配置是否匹配
    if config.get('model_name') == model_name:
        model_config = {k: v for k, v in config.items() if not isinstance(v, dict)}
        print(f"使用基础配置进行推理")
    else:
        # 在实验配置中查找
        for exp_name, exp_config in config.items():
            if isinstance(exp_config, dict) and exp_config.get('model_name') == model_name:
                model_experiment = exp_name
                model_config = get_experiment_config(config, exp_name)
                print(f"找到实验配置: {exp_name}")
                break
    
    if model_config is None:
        print(f"错误: 找不到名为 {model_name} 的模型配置")
        return None
    
    # 确保模型目录存在
    os.makedirs(model_config['model_path'], exist_ok=True)
    
    # 设置随机种子
    same_seeds(model_config.get('seed', 0))
    
    # 设置设备
    try:
        paddle.device.set_device('gpu:0' if model_config.get('use_gpu', False) else 'cpu')
    except:
        print("GPU不可用，使用CPU进行推理")
        paddle.device.set_device('cpu')
    
    # 确定数据路径
    if data_path is None:
        data_path = model_config.get('private_test_path')
        if data_path is None:
            print("错误: 未指定数据路径，且配置中没有private_test_path")
            return None
    
    # 确定输出文件
    if output_file is None:
        output_file = model_config.get('prediction_file')
        if output_file is None:
            output_file = "predictions.csv"
            print(f"未指定输出文件，使用默认文件: {output_file}")
    
    # 加载数据
    try:
        print(f"加载数据: {data_path}")
        test_data = np.load(data_path)
        print(f"数据形状: {test_data.shape}")
        
        # 创建数据集和数据加载器
        test_dataset = TIMITDataset(test_data)
        test_loader = paddle.io.DataLoader(
            test_dataset, 
            batch_size=batch_size, 
            shuffle=False
        )
        
        # 创建模型
        model = Classifier(model_config)
        
        # 加载模型参数
        model_path = f"{model_config['model_path']}/{model_name}_best_model.pdparams"
        if not os.path.exists(model_path):
            print(f"错误: 找不到模型文件 {model_path}")
            return None
        
        print(f"加载模型参数: {model_path}")
        model_state_dict = paddle.load(model_path)
        model.set_state_dict(model_state_dict)
        
        # 进行推理
        model.eval()
        predictions = []
        
        print("开始推理...")
        for batch_id, data in enumerate(test_loader):
            x_data = data[0]
            logits = model(x_data)
            predictions.extend(paddle.argmax(logits, axis=1).cpu().numpy().tolist())
            if batch_id % 10 == 0:
                print(f"已处理 {batch_id * batch_size} 个样本")
        
        # 保存预测结果
        with open(output_file, "w") as f:
            for pred in predictions:
                f.write(f"{pred}\n")
        
        print(f"预测结果已保存到 {output_file}")
        return predictions
    
    except Exception as e:
        print(f"推理过程中出错: {e}")
        import traceback
        traceback.print_exc()
        return None

def run_all_experiments(config_file="config.yaml"):
    """逐一运行所有实验"""
    # 运行基线实验
    run_experiment(config_file=config_file)
    
    # 运行其他实验
    for experiment in EXPERIMENTS:
        run_experiment(experiment, config_file)
    
    # 汇总结果
    generate_report()

def generate_report():
    """从experiment_results.json生成实验报告"""
    try:
        with open('experiment_results.json', 'r', encoding='utf-8') as f:
            results = json.load(f)
        
        # 创建结果表格
        data = []
        for result in results:
            data.append({
                'Method': result['model_name'],
                'Description': result['description'],
                'Validation Score (Accuracy)': result['val_acc'],
                'Test Score (Accuracy)': result['test_acc']
            })
        
        df = pd.DataFrame(data)
        
        # 对结果进行排序
        df = df.sort_values(by='Test Score (Accuracy)', ascending=False)
        
        # 格式化结果
        df['Validation Score (Accuracy)'] = df['Validation Score (Accuracy)'].map(lambda x: f'{x:.4f}')
        df['Test Score (Accuracy)'] = df['Test Score (Accuracy)'].map(lambda x: f'{x:.4f}')
        
        # 保存报告
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        report_file = f'experiment_report_{timestamp}.md'
        
        with open(report_file, 'w', encoding='utf-8') as f:
            f.write("# TIMIT Phoneme Classification 实验报告\n\n")
            f.write("## 实验结果摘要\n\n")
            f.write("以下是各种模型配置的实验结果：\n\n")
            
            # 创建简洁的Markdown表格，避免使用to_markdown可能导致的编码问题
            f.write("| Method | Description | Validation Score (Accuracy) | Test Score (Accuracy) |\n")
            f.write("|--------|-------------|----------------------------|----------------------|\n")
            
            for _, row in df.iterrows():
                f.write(f"| {row['Method']} | {row['Description']} | {row['Validation Score (Accuracy)']} | {row['Test Score (Accuracy)']} |\n")
            
            f.write("\n\n## 最佳模型\n\n")
            
            # 获取最佳模型
            best_model = df.iloc[0]
            f.write(f"最佳模型: **{best_model['Method']}**\n\n")
            f.write(f"- 描述: {best_model['Description']}\n")
            f.write(f"- 验证集准确率: {best_model['Validation Score (Accuracy)']}\n")
            f.write(f"- 测试集准确率: {best_model['Test Score (Accuracy)']}\n\n")
            
            f.write("## 实验方法说明\n\n")
            
            for result in results:
                f.write(f"### {result['model_name']}\n\n")
                f.write(f"- 描述: {result['description']}\n")
                f.write(f"- 网络结构: {result['config'].get('layer_sizes', 'N/A')}\n")
                f.write(f"- 激活函数: {result['config'].get('activation', 'N/A')}\n")
                f.write(f"- 批归一化: {result['config'].get('use_bn', False)}\n")
                f.write(f"- Dropout: {result['config'].get('use_dropout', False)}\n")
                if result['config'].get('use_dropout', False):
                    f.write(f"- Dropout率: {result['config'].get('dropout_rate', 0.2)}\n")
                f.write(f"- L2正则化: {result['config'].get('use_l2_reg', False)}\n")
                if result['config'].get('use_l2_reg', False):
                    f.write(f"- 权重衰减系数: {result['config'].get('weight_decay', 0.0001)}\n")
                f.write(f"- 最后一层Softmax: {result['config'].get('use_last_layer_softmax', False)}\n")
                f.write(f"- 学习率: {result['config'].get('learning_rate', 'N/A')}\n")
                f.write(f"- 训练轮数: {result['config'].get('epochs', 'N/A')}\n")
                f.write(f"- 批大小: {result['config'].get('batch_size', 'N/A')}\n\n")
            
            f.write("\n\n## 结论\n\n")
            f.write("通过以上实验，我们可以得出以下结论：\n\n")
            f.write("1. 深度网络（4层）比浅层网络（1层）表现更好\n")
            f.write("2. ReLU激活函数比Sigmoid激活函数更适合此任务\n")
            f.write("3. 批归一化有助于提高模型性能\n")
            f.write("4. 增加训练轮数可以进一步提高模型性能\n")
            f.write("5. 添加Dropout可以防止过拟合，提高模型泛化能力\n")
            f.write("6. L2正则化可以减轻模型过拟合，提高泛化能力\n")
            f.write("7. 在隐藏层使用ReLU，而在最后一层使用Softmax可以改变模型的概率输出特性\n")
            f.write("8. 更大的batch size（从64增加到256）可以加速训练，但对模型性能的影响需要具体分析\n")
            
        print(f"实验报告已保存到 {report_file}")
        return report_file
    except Exception as e:
        print(f"生成报告时出错: {e}")
        return None

def main():
    parser = argparse.ArgumentParser(description='运行TIMIT音素分类实验')
    parser.add_argument('--config', type=str, default='config.yaml', help='配置文件路径')
    parser.add_argument('--report', action='store_true', help='直接生成报告（不运行实验）')
    parser.add_argument('--experiment', type=str, help='运行特定实验')
    parser.add_argument('--all', action='store_true', help='运行所有实验')
    parser.add_argument('--inference', action='store_true', help='使用训练好的模型进行推理')
    parser.add_argument('--model', type=str, help='用于推理的模型名称')
    parser.add_argument('--data', type=str, help='推理数据集路径')
    parser.add_argument('--output', type=str, help='推理结果输出文件路径')
    parser.add_argument('--batch_size', type=int, default=64, help='推理批处理大小')
    
    args = parser.parse_args()
    
    if args.inference:
        if not args.model:
            print("错误: 进行推理需要指定模型名称，请使用 --model 参数")
            return
        inference_with_model(
            model_name=args.model, 
            data_path=args.data, 
            config_file=args.config, 
            batch_size=args.batch_size, 
            output_file=args.output
        )
    elif args.report:
        # 只生成报告，不运行实验
        generate_report()
    elif args.experiment:
        # 运行特定实验
        run_experiment(args.experiment, args.config)
    elif args.all:
        # 运行所有实验
        run_all_experiments(args.config)
    else:
        # 直接调用a1.py --all参数
        subprocess.run(["python", "a1.py", "--config", args.config, "--all"])
        generate_report()

if __name__ == '__main__':
    main() 